"""
场景需求
员工提交采购申请（包含物品、金额、原因）；
第一步：调用 “预算查询工具” 检查部门预算是否充足；
第二步：若金额≤5000 元，直接部门经理审批；若金额 > 5000 元，需额外 CEO 审批；
第三步：审批通过后，调用 “财务打款工具” 执行付款；
全程记录流程状态（用记忆组件保存），支持查询当前进度。
"""
from datetime import datetime

from langchain.agents import ConversationalAgent, AgentExecutor, StructuredChatAgent
from langchain.memory import ConversationBufferMemory
from langchain_core.messages import get_buffer_string
from langchain_core.prompts import PromptTemplate
from langchain_core.tools import tool
from langchain_openai import ChatOpenAI


@tool(
    description="用于查询部门预算是否充足。参数:department(必填,部门,示例:技术部,市场部,运营部),amount(必填,金额,示例:10000),返回值:部门预算充足/不充足,部门剩余预算 10000 元。本次申请金额 10000 元"
)
def check_budget_tool(department: str, amount: float) -> str:
    budget_data = {
        "技术部": 10000,
        "运营部": 5009,
        "市场部": 20000,
    }
    remaining = budget_data.get(department, 0)
    if remaining >= amount:
        return f"部门预算充足,{department} 部门剩余预算 {remaining} 元。本次申请金额 {amount} 元"

    else:
        return f"部门预算不足,{department} 部门剩余预算 {remaining} 元。本次申请金额 {amount} 元"


@tool(
    description="用于采购申请审批流程。参数:approver(必填,审批人,示例:部门经理,首席执行官,财务总监),content(必填,申请内容),amount(必填,申请金额),返回值:审批结果和下一步操作建议"
)
def purchase_approval_tool(approver: str, content: str, amount: float) -> str:
    print(f"agent_flow.py:purchase_approval_tool:approver={approver}, amount={amount}")

    # 审批逻辑
    if "部门经理" in approver:
        if amount <= 5000:
            return f"部门经理审批通过,审批内容:{content}。金额≤5000元，审批流程完成，可以进入付款阶段。"
        else:
            return f"部门经理审批通过,审批内容:{content}。金额>5000元，需要首席执行官进一步审批。"

    elif "首席执行官" in approver:
        return f"首席执行官审批通过,审批内容:{content}。所有审批完成，可以进入付款阶段。"

    elif "财务总监" in approver:
        return f"财务总监审批通过,审批内容:{content}。财务审批完成，可以进入付款阶段。"

    else:
        return f"审批不通过,审批人:{approver}无权限审批此申请,审批内容:{content}"


@tool(
    description="用于执行付款。参数:receiver(必填,收款人,示例:张三,王五),amount(必填,金额,示例:10000),返回值:付款成功,付款人:张三,金额:10000元,交易单号:PAY1234567890,交易时间:2023-01-01 00:00:00"
)
def finance_payment_tool(receiver: str, amount: float) -> str:
    return f"财务已向{receiver}打款{amount}元，交易单号：PAY{hash(receiver + str(amount))},交易时间：{datetime.now()}"


@tool(description="用于查询当前流程状态。参数:申请ID:当前步骤、处理人、结果")
def query_process_status_tool(apply_id: str) -> str:
    history = memory.load_memory_variables({})["apply_chat_history"]
    print("history:", history)
    history_str = get_buffer_string(history)
    return f"""
    申请ID:{apply_id},
    当前的状态:{history_str}
    """


tools = [
    check_budget_tool,
    purchase_approval_tool,
    finance_payment_tool,
    query_process_status_tool
]

llm = ChatOpenAI(
    model_name="qwen-plus",
    api_key="sk-6434b554122c4b049ceb805d703f695b",
    base_url="https://dashscope.aliyuncs.com/compatible-mode/v1",
)

# 核心:创建会话记忆,
memory = ConversationBufferMemory(
    memory_key="apply_chat_history",
    return_messages=True,  # 返回Message对象，方便Agent理解
)

prompt_template = """
你是一个采购流程审批的智能助手，你拥有以下工具:
{tools}

使用规则:
1:必须遵守以下工作流程的顺序, 禁止跳步:
    -第一步[必须执行]:收到采购申请后,无论什么情况, 都要先调用check_budget_tool检查部门预算
    -第二步[只有预算充足的情况才执行]:调用purchase_approval_tool进行审批
        -若金额≤5000元:approver="部门经理",审批完成后直接进入付款流程
        -若金额>5000元:先approver="部门经理",审批通过后再approver="首席执行官",两次审批都通过才能进入付款流程
    -第三步[所有审批申请通过才执行]:调用finance_payment_tool向申请人执行付款;
    
2: 执行query_process_status_tool,可以查询审批的状态 ,只能用户主动查询    

3:对话历史:{apply_chat_history}  

4:用户的采购申请:{input}

重要提醒:
- 对于金额>5000元的申请，必须看到两次审批通过的结果才能执行付款
- 审批工具会返回下一步操作建议，请严格按照建议执行
- 只有看到"可以进入付款阶段"的提示才能调用付款工具

"""

prompt = PromptTemplate(
    template=prompt_template,
    input_variables=["tools", "apply_chat_history", "input"],
)

agent = StructuredChatAgent.from_llm_and_tools(
    llm=llm,
    tools=tools,
    prompt=prompt,
    verbose=True
)

agent_executor = AgentExecutor(
    agent=agent,
    tools=tools,
    verbose=True,
    handle_parsing_errors=True,
    max_iterations=5,
    memory=memory,  # 关键:把记忆绑定到执行器
)

if __name__ == '__main__':
    res1 = agent_executor.invoke(
        {"input": "申请ID:YY001,申请人:杨涛,部门:技术部,采购物品:服务器,金额:8000,原因:服务器不够用了"})
    print(res1["output"])

    res2 = agent_executor.invoke({"input": "申请ID:YY001查询当前进度"})
    print(res2["output"])
